import numpy as np
import pandas as pd
import os
import time

class Feature_calc_df():
    root = r"data/stock_data/daily/consentrate_price"
    # root = r'data\etf_data\consentrate_daily_price'
    close = pd.read_pickle(os.path.join(root, "close.pkl.gzip"))
    high = pd.read_pickle(os.path.join(root, "high.pkl.gzip"))
    low = pd.read_pickle(os.path.join(root, "low.pkl.gzip"))
    lagRet = pd.read_pickle(os.path.join(root, "lagret.pkl.gzip"))
    volume = pd.read_pickle(os.path.join(root, "volume.pkl.gzip"))
    codes = (lagRet).columns
    org_shape = (lagRet).shape 
    datetime_list = pd.to_datetime(lagRet.index)
    
    def __init__(self, root,N, is_reset_daily = False):
        if not os.path.exists(root):
            os.makedirs(root)
        self.root = root
        self.is_reset_daily = is_reset_daily
        self.N = N
        self.s_N = int(0.5 * self.N)
        self.l_N = int(self.N)  
        self.vol_num = 500    

    def break_diff(self):           
        high =  Feature_calc_df.high
        low =  Feature_calc_df.low
        if self.is_reset_daily:
            h_mean = high.groupby(high.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).mean()))
            l_mean = low.groupby(low.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).mean()))
        else:
            h_mean = high.rolling(self.N, min_periods=1).mean()
            l_mean = low.rolling(self.N, min_periods=1).mean()
            
        high_ratio = (high >= h_mean ) * (np.divide((high - h_mean), h_mean))
        low_ratio = (low <= l_mean ) * (np.divide(( l_mean - low), l_mean))     
        f = high_ratio - low_ratio
        return f

    def break_diff_ma(self):   
        high =  Feature_calc_df.high
        low =  Feature_calc_df.low
       
        if self.is_reset_daily:     
            h_mean_s_N = high.groupby(high.index.date,as_index=False ).apply(lambda x: (x.rolling(self.s_N,  min_periods=1).mean()))
            h_mean_l_N = high.groupby(high.index.date,as_index=False ).apply(lambda x: (x.rolling(self.l_N,  min_periods=1).mean()))
            l_mean_s_N = low.groupby(low.index.date,as_index=False ).apply(lambda x: (x.rolling(self.s_N,  min_periods=1).mean()))
            l_mean_l_N = low.groupby(low.index.date,as_index=False ).apply(lambda x: (x.rolling(self.l_N,  min_periods=1).mean())) 
        else:  
            h_mean_s_N = high.rolling(self.s_N,  min_periods=1).mean()
            h_mean_l_N = high.rolling(self.l_N,  min_periods=1).mean()
            l_mean_s_N = low.rolling(self.s_N,  min_periods=1).mean()
            l_mean_l_N = low.rolling(self.l_N,  min_periods=1).mean()
            
        high_ratio = (h_mean_s_N >= h_mean_l_N ) * (np.divide((h_mean_s_N - h_mean_l_N), h_mean_l_N))
        low_ratio = (l_mean_s_N <= l_mean_l_N ) * (np.divide(( l_mean_l_N - l_mean_s_N), l_mean_l_N))     
        f = high_ratio - low_ratio
        return f
    
    def DS(self):     #试算查看2018年8，9月情况
        lagRet_normal = Feature_calc_df.lagRet
        lagRet_normal_df = lagRet_normal
        lagRet_abs_df = np.abs(lagRet_normal) 

        if self.is_reset_daily:          
            lagRet_normal_df = lagRet_normal_df.groupby(lagRet_normal_df.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).sum()))
            lagRet_abs_df = lagRet_abs_df.groupby(lagRet_abs_df.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).sum()))
        else:
            lagRet_normal_df = lagRet_normal_df.rolling(self.N, min_periods=1).sum()
            lagRet_abs_df = lagRet_abs_df.rolling(self.N, min_periods=1).sum()
            
        f = np.divide(lagRet_normal_df, lagRet_abs_df)
        return f 
    
    def lagRet_R(self):          
        lagRet =  Feature_calc_df.lagRet
        if self.is_reset_daily:
            lagRet_r = lagRet.groupby(lagRet.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).sum()))
        else:
            lagRet_r = lagRet.rolling(self.N, min_periods=1).sum()
        f = lagRet_r
        return f  
    
    def p_loc(self):    
        close =  Feature_calc_df.close
        high =  Feature_calc_df.high
        low =  Feature_calc_df.low
              
        if self.is_reset_daily:       
            high_max = high.groupby(high.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).max()))
            low_min = low.groupby(low.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).min()))    
        else:
            high_max = high.rolling(self.N, min_periods=1).max()
            low_min = low.rolling(self.N, min_periods=1).min()
        f = np.divide(close-low_min, high_max-low_min)
        return f
 
    def trading_SMA_Mmin(self):
        close =  Feature_calc_df.close
        if self.is_reset_daily:
            close_s = close.groupby(close.index.date,as_index=False ).apply(lambda x: (x.rolling(self.s_N,  min_periods=1).mean()))
            close_l = close.groupby(close.index.date,as_index=False ).apply(lambda x: (x.rolling(self.l_N,  min_periods=1).mean()))
        else:
            close_s = close.rolling(self.s_N,  min_periods=1).mean()
            close_l = close.rolling(self.l_N,  min_periods=1).mean()
        f = np.log(close_s/close_l)
        return f

    def std_rise_s_l(self):  
        lagret = Feature_calc_df.lagRet
        if self.is_reset_daily:
            std_s = lagret.groupby(lagret.index.date,as_index=False ).apply(lambda x: (x.rolling(self.s_N,  min_periods=1).std()))
            std_l = lagret.groupby(lagret.index.date,as_index=False ).apply(lambda x: (x.rolling(self.l_N,  min_periods=1).std()))
        else:
            std_s = lagret.rolling(self.s_N,  min_periods=1).std()
            std_l = lagret.rolling(self.l_N,  min_periods=1).std()
        f = (std_s / std_l.values -1)
        return f

    def return_stasistic(self, target):#特殊
        lagRet_vol = Feature_calc_df.lagRet
        if self.is_reset_daily:
            if  target == "max":
                f = lagRet_vol.groupby(lagRet_vol.index.date,as_index=False).apply(lambda x: x.rolling(self.N, min_periods=1).max()) 
            if  target == "min":
                f = lagRet_vol.groupby(lagRet_vol.index.date,as_index=False).apply(lambda x: x.rolling(self.N, min_periods=1).min()) 
            if  target == "mean":
                f = lagRet_vol.groupby(lagRet_vol.index.date,as_index=False).apply(lambda x: x.rolling(self.N, min_periods=1).mean()) 
            if  target == "median":
                f = lagRet_vol.groupby(lagRet_vol.index.date,as_index=False).apply(lambda x: x.rolling(self.N, min_periods=1).median()) 
            if  target == "skew":
                f = lagRet_vol.groupby(lagRet_vol.index.date,as_index=False).apply(lambda x: x.rolling(self.N, min_periods=1).skew()) 
            if  target == "kurt":
                f = lagRet_vol.groupby(lagRet_vol.index.date,as_index=False).apply(lambda x: x.rolling(self.N, min_periods=1).kurt())
        else:
            if  target == "max":
                f = lagRet_vol.rolling(self.N, min_periods=1).max()
            if  target == "min":
                f = lagRet_vol.rolling(self.N, min_periods=1).min()
            if  target == "mean":
                f = lagRet_vol.rolling(self.N, min_periods=1).mean()
            if  target == "median":
                f = lagRet_vol.rolling(self.N, min_periods=1).median()
            if  target == "skew":
                f = lagRet_vol.rolling(self.N, min_periods=1).skew()
            if  target == "kurt":
                f = lagRet_vol.rolling(self.N, min_periods=1).kurt()
        f = f[f.index>'2005-01-01']
        f.to_pickle(f"{self.root}/return_stasistic_{self.N}_{ target}.pkl.gzip")      

    def return_sign(self):
        lagRet_df = Feature_calc_df.lagRet
        if self.is_reset_daily:        
            f = np.sign(lagRet_df).groupby(lagRet_df.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).sum()))
        else:
            f = np.sign(lagRet_df).rolling(self.N, min_periods=1).sum()
        f = f/self.N
        return f
    
    def f5(self):
        high =  Feature_calc_df.high
        low =  Feature_calc_df.low
        close_df =  Feature_calc_df.close
        if self.is_reset_daily:
            h_t = high.groupby(high.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).max()).shift(1))
            l_t = low.groupby(low.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).min()).shift(1))
            m_t = (h_t + l_t + 2 * close_df.groupby(close_df.index.date,as_index=False ).apply(lambda x: x.shift(1))) / 4
            p_t = close_df - m_t
            f_numerator = p_t.groupby(p_t.index.date,as_index=False).apply(lambda x: (x.rolling(int(0.5*self.N)).mean()))         
        else:
            h_t = high.rolling(self.N, min_periods=1).max().shift(1)
            l_t = low.rolling(self.N, min_periods=1).min().shift(1)
            m_t = (h_t + l_t + 2 * close_df.shift(1)) / 4
            p_t = close_df - m_t
            f_numerator = p_t.rolling(int(0.5*self.N)).mean()
        f = f_numerator/close_df
        return f

    def return_vwap(self):
        lagRet_v = Feature_calc_df.lagRet * Feature_calc_df.volume    
        lagRet_v_df =  lagRet_v
        volume_df =  Feature_calc_df.volume
        if self.is_reset_daily:
            cumret_N = lagRet_v_df.groupby(lagRet_v_df.index.date,as_index=False ).apply(lambda x:  (x.rolling(self.N, min_periods=1).sum()))
            cumVolume_N = volume_df.groupby(volume_df.index.date,as_index=False ).apply(lambda x:  (x.rolling(self.N, min_periods=1).sum()))
        else:
            cumret_N = lagRet_v_df.rolling(self.N, min_periods=1).sum()
            cumVolume_N = volume_df.rolling(self.N, min_periods=1).sum()
        f = cumret_N / cumVolume_N
        f = cut_tail(f)
        return f
  
    def vwap_loc(self):
        close_df = Feature_calc_df.close
        volume_df = Feature_calc_df.volume
        close_v_df = close_df * volume_df
        if self.is_reset_daily:
            close_min = close_df.groupby(close_df.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).min()))
            close_max = close_df.groupby(close_df.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).max()))                 
            vwap_numerator = close_v_df.groupby(close_v_df.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).sum()))
            vwap_denominnator = volume_df.groupby(volume_df.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).sum()))
        else:
            close_min = close_df.rolling(self.N, min_periods=1).min()
            close_max = close_df.rolling(self.N, min_periods=1).max()                   
            vwap_numerator = close_v_df.rolling(self.N, min_periods=1).sum()
            vwap_denominnator = volume_df.rolling(self.N, min_periods=1).sum()
        vwap = np.divide(vwap_numerator, vwap_denominnator)
        f = (vwap-close_min)/(close_max-close_min)
        return f
        
    def vwap_s_l(self):
        close_df = Feature_calc_df.close
        volume_df = Feature_calc_df.volume
        close_v_df = close_df * volume_df

        if self.is_reset_daily: 
            
            vwap_s_numerator = close_v_df.groupby(close_v_df.index.date,as_index=False ).apply(lambda x: (x.rolling(int(self.N/2)).sum()))
            vwap_s_denominnator = volume_df.groupby(volume_df.index.date,as_index=False ).apply(lambda x: (x.rolling(int(self.N/2)).sum()))
            
            vwap_l_numerator = close_v_df.groupby(close_v_df.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).sum()))
            vwap_l_denominnator = volume_df.groupby(volume_df.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).sum()))

        else:

            vwap_s_numerator = close_v_df.rolling(int(self.N/2)).sum()
            vwap_s_denominnator = volume_df.rolling(int(self.N/2)).sum()
            
            vwap_l_numerator = close_v_df.rolling(self.N, min_periods=1).sum()
            vwap_l_denominnator = volume_df.rolling(self.N, min_periods=1).sum()

        vwap_s = np.divide(vwap_s_numerator, vwap_s_denominnator)  
        vwap_l = np.divide(vwap_l_numerator, vwap_l_denominnator)
        f = vwap_s/ vwap_l-1
        return f
  
    def deviation(self):
        self.a = 1.5
        close =  Feature_calc_df.close
        high =  Feature_calc_df.high
        low =  Feature_calc_df.low
        if self.is_reset_daily: 
            c_std = close.groupby(close.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).std()))
            c_h = high.groupby(high.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).max()))
            c_l = low.groupby(low.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).min())) 
            c_0 = close.groupby(close.index.date,as_index=False ).apply(lambda x:(x.shift(self.N)))     
        else:
            c_std = close.rolling(self.N, min_periods=1).std()
            c_h = high.rolling(self.N, min_periods=1).max()
            c_l = low.rolling(self.N, min_periods=1).min()  
            c_0 = close.shift(self.N)
            
        condition1 = (np.minimum(c_h.values-c_0.values, c_h.values-close.values) > (self.a * c_std.values)).astype(int)
        condition2 = (np.minimum(c_0.values-c_l.values, close.values-c_l.values) > (self.a * c_std.values)).astype(int)   
        # print("同时成立:",((condition1+condition2)>=2).any()) 
        f = (condition2 - condition1).reshape(Feature_calc_df.org_shape) 
        f = pd.DataFrame(f, index=close.index, columns=close.columns)

        return f
  
    def new_p_loc(self): 
        close = Feature_calc_df.close
        high = Feature_calc_df.high
        low = Feature_calc_df.low
              
        if self.is_reset_daily:   
            h_max = high.groupby(high.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).max()))
            l_min = low.groupby(low.index.date,as_index=False ).apply(lambda x: (x.rolling(self.N, min_periods=1).min()))  
            m = (h_max + l_min + 2*close)/4
            p = close - m       
            f = np.divide(p.groupby(p.index.date,as_index=False ).apply(lambda x: (x.rolling(self.s_N, min_periods=1).mean())), p.groupby(p.index.date,as_index=False ).apply(lambda x: (x.rolling(self.l_N,  min_periods=1).mean())) )  
        else:
            h_max = high.rolling(self.N, min_periods=1).max()
            l_min = low.rolling(self.N,  min_periods=1).min()
            m = (h_max + l_min + 2*close)/4
            p = close - m   
            f = (p/p.rolling(self.vol_num, min_periods=1).std())
        return f         
       
    def calc_run(self, func): 
        t1 = time.time()       
        signal_df = func()
        # signal_df = signal_df/signal_df.rolling(self.vol_num, min_periods=1).std()
        signal_df = signal_df.replace([np.inf, -np.inf], np.nan)
        signal_df = signal_df[signal_df.index>'2005-01-01']
        signal_df.to_pickle(f"{self.root}/{func.__name__}_{self.N}.pkl.gzip")
        print(f"{func.__name__}_{self.N}计算完毕!耗时{(time.time()-t1):.3f}s")

    def forward(self):
        self.calc_run(self.break_diff)
        self.calc_run(self.break_diff_ma)
        self.calc_run(self.DS)
        self.calc_run(self.p_loc)
        self.calc_run(self.lagRet_R)
        self.calc_run(self.trading_SMA_Mmin)
        self.calc_run(self.deviation)
        self.calc_run(self.std_rise_s_l)
        self.calc_run(self.return_vwap)
        self.calc_run(self.return_sign)
        self.calc_run(self.f5)
        self.calc_run(self.vwap_s_l)
        self.calc_run(self.vwap_loc)     
        target_lst = ["max", "min", "mean", "median", "kurt", "skew"]
        for i in target_lst:
            self.return_stasistic(target=i)      
              
    
def cut_tail(df, scale=1):
    df = scale*df 
    df[df>=1] = 1
    df[df<=-1] = -1
    return df

def cut_tail_element(data, scale=1):
    data = scale*data 
    data[data>=1] = 1
    data[data<=-1] = -1
    return data


